# -*- coding: utf-8 -*-
#
#
# Author: Alex Cai
# Email: cyy0523xc@qq.com
# Created Time: 2025-06-13
from langgraph.graph import END, StateGraph
from typing import Dict, TypedDict


# 定义状态结构
class AgentState(TypedDict):
    a: int
    b: int
    operation: str
    result: int


# 模拟大模型接口 - 决策操作类型
def mock_llm_decision(state: AgentState) -> Dict[str, str]:
    """模拟大模型决定操作类型（实际使用会替换为真实LLM调用）"""
    # 这里用简单逻辑模拟：正数加法，负数减法
    operation = "add" if state["a"] >= 0 and state["b"] >= 0 else "subtract"
    print(f"📡 模拟大模型决策: 建议操作 {operation.upper()}")
    return {"operation": operation}


# 计算节点函数
def add_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] + state["b"]
    print(f"➕ 加法结果: {state['a']} + {state['b']} = {result}")
    return {"result": result}


def subtract_numbers(state: AgentState) -> Dict[str, int]:
    result = state["a"] - state["b"]
    print(f"➖ 减法结果: {state['a']} - {state['b']} = {result}")
    return {"result": result}


# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点（实际应用中决策节点会调用真实LLM）
    workflow.add_node("model_decision", mock_llm_decision)
    workflow.add_node("add", add_numbers)
    workflow.add_node("subtract", subtract_numbers)

    # 设置入口点
    workflow.set_entry_point("model_decision")

    # 根据决策结果路由
    workflow.add_conditional_edges(
        "model_decision",
        lambda state: state["operation"],
        {
            "add": "add",
            "subtract": "subtract",
        }
    )

    # 连接终结节点
    workflow.add_edge("add", END)
    workflow.add_edge("subtract", END)

    return workflow.compile()


# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例1：正数触发加法
    print("\n测试案例 1: (5, 3)")
    result = app.invoke({"a": 5, "b": 3})
    print(f"最终结果: {result['result']}")

    # 测试案例2：负数触发减法
    print("\n测试案例 2: (10, -2)")
    result = app.invoke({"a": 10, "b": -2})
    print(f"最终结果: {result['result']}")

    # 测试案例3：自定义操作
    print("\n测试案例 3: (8, 4) 强制减法")
    result = app.invoke({"a": 8, "b": 4, "operation": "subtract"})
    print(f"最终结果: {result['result']}")
